package kibahed.uni.mdp.calculation;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public class MDP_Data {

	private int rows;
	private int collumns;
	private double highNoise;
	private double lowNoise;
	private double discount;
	private Map<Integer, String> actions;
	private Map<Integer, Integer> actionHelper;
	private double[][] rewardsPerActionAndState;
	private double[][] probabilitiesPerStateAndAction;
	private List<Integer> noisyStates;
	private List<Integer> goodStates;
	private List<Integer> badStates;
	private double[] noises;
	private double[] rewards;

	private boolean fromVisu;

	public MDP_Data(int rows, int collums, double lowNoise, double highNoise,
			double discount, boolean fromVisu) {
		this.fromVisu = fromVisu;
		this.rows = rows;
		this.collumns = collums;
		this.lowNoise = lowNoise;
		this.highNoise = highNoise;
		this.discount = discount;
		noises = new double[rows * collums];
		rewardsPerActionAndState = new double[rows * collums][5];
		createActions();
		createProbabilities();
		createRewardsPerActionAndState();
	}

	public MDP_Data(int rows, int collums, double discount, double[] noises,
			double[] rewards, boolean fromVisu) {
		this.fromVisu = fromVisu;
		this.rows = rows;
		this.rewards = rewards;
		this.noises = noises;
		this.collumns = collums;
		this.discount = discount;
		rewardsPerActionAndState = new double[rows * collums][5];
		createActions();
		createProbabilities();
		createRewardsPerActionAndState();
	}

	private void createActions() {
		actions = new HashMap<Integer, String>();
		actions.put(1, "up");
		actions.put(2, "down");
		actions.put(3, "left");
		actions.put(4, "right");
		actions.put(5, "stay");

		actionHelper = new HashMap<Integer, Integer>();
		actionHelper.put(1, -1);
		actionHelper.put(2, +1);
		actionHelper.put(3, -1);
		actionHelper.put(4, +1);
		actionHelper.put(5, 0);
	}

	public void setNoieses(double[] noises) {
		this.noises = noises;
	}

	public void setRewards(double[] rewards) {
		this.rewards = rewards;
	}

	private void createProbabilities() {
		noisyStates = new ArrayList<Integer>();
		noisyStates.add(6);
		noisyStates.add(7);
		noisyStates.add(8);
		noisyStates.add(12);
		probabilitiesPerStateAndAction = new double[rows * collumns][5];

		for (int i = 0; i < rows * collumns; i++) {
			double[] probs = new double[5];

			if (fromVisu) {
				for (int j = 0; j < probs.length; j++) {
					probs[j] = 1 - noises[i];
				}
			} else {
				if (noisyStates.contains(i)) {
					for (int j = 0; j < probs.length; j++) {
						probs[j] = 1 - highNoise;
					}
				} else {
					for (int j = 0; j < probs.length; j++) {
						probs[j] = 1 - lowNoise;
					}
				}
			}
			probabilitiesPerStateAndAction[i] = probs;

		}
	}

	private void createRewardsPerActionAndState() {

		// set positions of 'good' and 'bad' states
		goodStates = new ArrayList<Integer>();
		goodStates.add(14);
		badStates = new ArrayList<Integer>();
		badStates.add(1);
		badStates.add(3);
		badStates.add(11);
		badStates.add(13);

		for (int i = 0; i < rows; i++) {
			for (int j = 0; j < collumns; j++) {

				int index = i * collumns + j;

				for (int a : actions.keySet()) {

					boolean staysIn = staysInBounds(i, j, a);

					if (!fromVisu) {
						if (badStates.contains(index)) {
							rewardsPerActionAndState[index][a - 1] = -1;
						} else if (goodStates.contains(index)) {
							rewardsPerActionAndState[index][a - 1] = 1;
						} else {
							rewardsPerActionAndState[index][a - 1] = 0;
						}
					}

					if (!staysIn) {
						rewardsPerActionAndState[index][a - 1] = -1;
					} else if (fromVisu) {
						rewardsPerActionAndState[index][a - 1] = rewards[index];
					}
				}
			}
		}
	}

	private boolean staysInBounds(int row, int collumn, int action) {

		int movement = actionHelper.get(action);

		if (action == 1 || action == 2) {

			int newRow = row + movement;
			if (newRow >= rows || newRow < 0) {
				return false;
			}

		} else if (action == 3 || action == 4) {

			int newCollumn = collumn + movement;
			if (newCollumn >= collumns || newCollumn < 0) {
				return false;
			}
		}

		return true;
	}

	public Map<Integer, Integer> getActionNeighbourMap(int state) {

		Map<Integer, Integer> neighbours = new HashMap<Integer, Integer>();

		for (int a = 1; a <= 5; a++) {
			neighbours.put(a, getNeighbourWithAction(state, a));
		}

		return neighbours;
	}

	private int getNeighbourWithAction(int state, int action) {

		int collumn = state % collumns;
		int row = state / rows;

		if (action == 1) {
			if ((row - 1) >= 0) {
				return state - collumns;
			}
		}

		if (action == 2) {
			if ((row + 1) <= 4) {
				return state + collumns;
			}
		}

		if (action == 3) {
			if ((collumn - 1) >= 0) {
				return state - 1;
			}
		}

		if (action == 4) {
			if ((collumn + 1) <= 4) {
				return state + 1;
			}
		}

		if (action == 5) {
			return state;
		}

		// returns -1 if MDP wants to leave the field
		return -1;
	}

	@SuppressWarnings("unused")
	private void printRewardsPerAction() {
		for (int i = 0; i < rows * collumns; i++) {
			System.out.println("\n-" + (i + 1)
					+ "---------------------------------------------------\n");
			for (int j = 0; j < 5; j++) {
				System.out.print("| " + rewardsPerActionAndState[i][j]);
			}
		}
	}

	public double getDiscount() {
		return discount;
	}

	public void setDiscount(double discount) {
		this.discount = discount;
	}

	public Map<Integer, String> getActions() {
		return actions;
	}

	public void setActions(Map<Integer, String> actions) {
		this.actions = actions;
	}

	public double[][] getRewardsPerActionAndState() {
		return rewardsPerActionAndState;
	}

	public void setRewardsPerActionAndState(double[][] rewardsPerActionAndState) {
		this.rewardsPerActionAndState = rewardsPerActionAndState;
	}

	public double[][] getProbabilitiesPerStateAndAction() {
		return probabilitiesPerStateAndAction;
	}

	public int getRows() {
		return rows;
	}

	public int getCollumns() {
		return collumns;
	}

}
